from espdl import HumanDetector
from jpeg import Decoder, Encoder


decoder = Decoder()
encoder = Encoder(width=640, height=480,pixel_format="RGB888")
human_detector = HumanDetector(width=640, height=480)

# 捕获并处理图像
img = open("pedestrian.jpg", "rb").read()  # 捕获原始图像（通常是JPEG格式）
framebuffer = decoder.decode(img)  # 转换为RGB888格式
# 将memoryview转换为bytearray以便修改
framebuffer = bytearray(framebuffer)
# 运行行人检测
results = human_detector.run(framebuffer)

# 绘制边框
def draw_rectangle(buffer, width, height, x, y, w, h, color=(255, 0, 0)):
    """
    在RGB888格式的图像缓冲区上绘制矩形边框
    :param buffer: 图像缓冲区
    :param width: 图像宽度
    :param height: 图像高度
    :param x: 边框左上角的x坐标
    :param y: 边框左上角的y坐标
    :param w: 边框宽度
    :param h: 边框高度
    :param color: 边框颜色（RGB格式）
    """
    # 辅助函数：设置单个像素的颜色
    def set_pixel(buffer, width, x, y, color):
        offset = (y * width + x) * 3
        buffer[offset] = color[0]  # R
        buffer[offset + 1] = color[1]  # G
        buffer[offset + 2] = color[2]  # B

    # 绘制上边框
    for i in range(x, x + w):
        if 0 <= i < width and 0 <= y < height:
            set_pixel(buffer, width, i, y, color)

    # 绘制下边框
    for i in range(x, x + w):
        if 0 <= i < width and 0 <= y + h < height:
            set_pixel(buffer, width, i, y + h, color)

    # 绘制左边框
    for j in range(y, y + h):
        if 0 <= j < height and 0 <= x < width:
            set_pixel(buffer, width, x, j, color)

    # 绘制右边框
    for j in range(y, y + h):
        if 0 <= j < height and 0 <= x + w < width:
            set_pixel(buffer, width, x + w, j, color)

# 在图像上绘制边框
for face in results:
    print(face)
    x1, y1, x2, y2 = face['box']
    draw_rectangle(framebuffer, 640, 480, x1, y1, x2-x1, y2-y1, color=(255, 0, 0))  # 使用红色边框

# 将带有边框的图像重新编码为JPEG格式并保存
marked_img = encoder.encode(framebuffer)
with open("pedestrian_marked.jpg", "wb") as f:
    f.write(marked_img)